#include <iostream>
#include <vector>

#include <map>
#include <set>

#include <cassert>

#include "sds/sfile.h"
#include "sds/sdataset.h"
#include "sds/sdataview.h"

#include <lemon/list_graph.h>
#include <lemon/graph_utils.h>
#include <lemon/bin_heap.h>
#include <lemon/unionfind.h>
#include <lemon/xy.h>

using namespace lemon;
using namespace std;
using namespace sds;

typedef ListUGraph Graph;
typedef Graph::Node Node;
typedef Graph::UEdge UEdge;
typedef Graph::NodeIt NodeIt;
typedef Graph::UEdgeIt UEdgeIt;
typedef Graph::IncEdgeIt IncEdgeIt;

int main(int argc, const char *argv[]) {

  assert(argc >= 3);

  SFile sourceFile(argv[1]);
  
  SDataView<unsigned> source = 
    sourceFile.getDataset(0).getDataView<unsigned>();
  vector<SInteger> size = source.getSize();
  
  Graph graph;
  map<unsigned, Node> nodes;

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (source[i][j] == 0) continue;
      if (nodes.find(source[i][j]) == nodes.end()) {
	Node node = graph.addNode();
	nodes.insert(make_pair(source[i][j], node));
      }
    }
  }

  cout << "Nodes found: " << countNodes(graph) << endl;

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (source[i][j] == 0) continue;
      if (j + 1 < size[1]) {
	if (source[i][j] != source[i][j + 1] && source[i][j + 1] != 0) {
	  if (findEdge(graph, nodes[source[i][j]], nodes[source[i][j + 1]]) 
	      == INVALID) {
	    graph.addEdge(nodes[source[i][j]], nodes[source[i][j + 1]]);
	  }
	}
      }
      if (i + 1 < size[0]) {
	if (source[i + 1][j] != 0 && source[i][j] != source[i + 1][j]) {
	  if (findEdge(graph, nodes[source[i][j]], nodes[source[i + 1][j]]) 
	      == INVALID) {
	    graph.addEdge(nodes[source[i][j]], nodes[source[i + 1][j]]);
	  }
	}
      }
    }
  }

  cout << "Edges found: " << countUEdges(graph) << endl;

  Graph::NodeMap<char> color(graph, -2);
  
  Graph::NodeMap<int> heapMap(graph, -1);
  BinHeap<Node, int, Graph::NodeMap<int> > heap(heapMap);
  
  for (NodeIt it(graph); it != INVALID; ++it) {
    heap.push(it, countOutEdges(graph, it));
  }

  vector<Node> order;

  while (!heap.empty()) {
    Node node = heap.top();
    heap.pop();
    color[node] = -1;
    order.push_back(node);
    for (IncEdgeIt it(graph, node); it != INVALID; ++it) {
      Node target = graph.runningNode(it); 
      if (color[target] == -2) {
	heap.decrease(target, heap[target] - 1);
      }
    }
  }

  for (int i = order.size() - 1; i >= 0; --i) {
    set<int> colors;
    for (IncEdgeIt it(graph, order[i]); it != INVALID; ++it) {
      Node target = graph.runningNode(it); 
      if (color[target] != -1) {
	colors.insert(color[target]);
      }
    }
    for (int j = 1; j <= 6; ++j) {
      if (colors.find(j) == colors.end()) {
	color[order[i]] = j;
	break;
      }
    }
  }

  SDataView<char> target(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (source[i][j] == 0) {
	target[i][j] = 0;
      } else {
	target[i][j] = color[nodes[source[i][j]]];
      }
    }
  }

  SFile targetFile(argv[2], SFile::CREATE);
  targetFile.createDataset("segmented", UINT8, size).
    setDataView<char>(target);

  return 0;
}
